import os
import time
import tushare as ts
import pandas as pd
from datetime import datetime, timedelta

# 尝试导入 TA-Lib；若不可用则使用 pandas 回退实现
try:
    import talib as ta  # type: ignore
    HAS_TALIB = True
except Exception:
    ta = None
    HAS_TALIB = False

# 1. 初始化Tushare（从环境变量或代码中提供 Token）
# 优先从环境变量 TUSHARE_TOKEN 读取，避免把 Token 写进代码仓库
TUSHARE_TOKEN = os.getenv("TUSHARE_TOKEN", "你的Tushare Token")
if not TUSHARE_TOKEN or TUSHARE_TOKEN.strip() == "" or TUSHARE_TOKEN == "你的Tushare Token":
    raise RuntimeError(
        "未检测到有效的 Tushare Token。请设置环境变量 TUSHARE_TOKEN 或在代码中替换为你的实际 Token。\n"
        "PowerShell(临时会话)示例: $env:TUSHARE_TOKEN=\"你的Token\""
    )
ts.set_token(TUSHARE_TOKEN)
pro = ts.pro_api()

# 2. 定义要获取的股票代码（2支示例）
stock_codes = ["000001.SZ", "600519.SH"]  # 平安银行、贵州茅台


# 3. 获取股票基本信息（对应stocks表）
def get_stock_basic():
    try:
        # 优先尝试一次性查询多代码（可减少接口调用次数）
        codes_str = ",".join(stock_codes)
        try:
            df = pro.stock_basic(ts_code=codes_str, fields="ts_code,symbol,name,area,industry,list_date")
            if df is not None and len(df) > 0:
                df["id"] = range(1, len(df) + 1)
                df = df[["id", "ts_code", "symbol", "name", "area", "industry", "list_date"]]
                return df
        except Exception:
            # 触发限频等异常时等待后重试一次
            time.sleep(61)
            df = pro.stock_basic(ts_code=codes_str, fields="ts_code,symbol,name,area,industry,list_date")
            if df is not None and len(df) > 0:
                df["id"] = range(1, len(df) + 1)
                df = df[["id", "ts_code", "symbol", "name", "area", "industry", "list_date"]]
                return df

        # 回退为逐个代码查询，并在调用间隙做限速等待
        parts = []
        for idx, code in enumerate(stock_codes):
            try:
                df_part = pro.stock_basic(ts_code=code, fields="ts_code,symbol,name,area,industry,list_date")
            except Exception:
                time.sleep(61)
                df_part = pro.stock_basic(ts_code=code, fields="ts_code,symbol,name,area,industry,list_date")
            if df_part is not None and len(df_part) > 0:
                parts.append(df_part)
            if idx < len(stock_codes) - 1:
                time.sleep(61)
        if parts:
            df = pd.concat(parts, ignore_index=True)
            df["id"] = range(1, len(df) + 1)
            df = df[["id", "ts_code", "symbol", "name", "area", "industry", "list_date"]]
            return df
    except Exception:
        # 最终回退：构造最小可用的基础信息表，避免脚本中断
        df = pd.DataFrame({
            "ts_code": stock_codes,
            "symbol": [c.split(".")[0] for c in stock_codes],
            "name": [None] * len(stock_codes),
            "area": [None] * len(stock_codes),
            "industry": [None] * len(stock_codes),
            "list_date": [None] * len(stock_codes),
        })
        df["id"] = range(1, len(df) + 1)
        df = df[["id", "ts_code", "symbol", "name", "area", "industry", "list_date"]]
        return df


# 4. 获取股票日线行情（对应stock_daily_data表）
def get_stock_daily():
    daily_data = []
    # 取最近30天的数据（避免数据量过大）
    end_date = datetime.now().strftime("%Y%m%d")
    start_date = (datetime.now() - timedelta(days=30)).strftime("%Y%m%d")
    
    for ts_code in stock_codes:
        df = pro.daily(
            ts_code=ts_code,
            start_date=start_date,
            end_date=end_date,
            fields="ts_code,trade_date,open,high,low,close,vol,amount"
        )
        if df is not None and len(df) > 0:
            daily_data.append(df)
        # 限速等待，避免免费账户一分钟多次调用受限
        if ts_code != stock_codes[-1]:
            time.sleep(61)
    
    # 合并数据并添加自增主键
    if daily_data:
        df_all = pd.concat(daily_data, ignore_index=True)
    else:
        df_all = pd.DataFrame(columns=["id","ts_code","trade_date","open","high","low","close","vol","amount"])
        return df_all

    # 按日期升序便于后续技术指标计算
    df_all = df_all.sort_values(["ts_code", "trade_date"]).reset_index(drop=True)

    # 转成数值类型（若API返回为字符串）
    for col in ["open", "high", "low", "close", "vol", "amount"]:
        df_all[col] = pd.to_numeric(df_all[col], errors="coerce")
    df_all["id"] = range(1, len(df_all)+1)
    # 调整字段顺序
    df_all = df_all[["id", "ts_code", "trade_date", "open", "high", "low", "close", "vol", "amount"]]
    return df_all


# 5. 计算技术指标（对应stock_technical_indicators表）
def get_technical_indicators(daily_df):
    tech_data = []
    for ts_code in stock_codes:
        # 筛选单支股票的日线数据（按日期升序排列）
        stock_daily = daily_df[daily_df["ts_code"] == ts_code].sort_values("trade_date")
        if stock_daily.empty:
            continue
        close_s = pd.to_numeric(stock_daily["close"], errors="coerce")
        
        # 计算技术指标
        if HAS_TALIB:
            ma5 = pd.Series(ta.SMA(close_s.values, timeperiod=5), index=close_s.index)
            ma20 = pd.Series(ta.SMA(close_s.values, timeperiod=20), index=close_s.index)
            rsi = pd.Series(ta.RSI(close_s.values, timeperiod=14), index=close_s.index)
            macd, macd_signal, macd_hist = ta.MACD(close_s.values)
            macd = pd.Series(macd, index=close_s.index)
            boll_upper, boll_mid, boll_lower = ta.BBANDS(close_s.values, timeperiod=20)
            boll_upper = pd.Series(boll_upper, index=close_s.index)
            boll_lower = pd.Series(boll_lower, index=close_s.index)
        else:
            # pandas 回退实现
            ma5 = close_s.rolling(window=5, min_periods=5).mean()
            ma20 = close_s.rolling(window=20, min_periods=20).mean()
            delta = close_s.diff()
            gain = delta.clip(lower=0)
            loss = -delta.clip(upper=0)
            avg_gain = gain.ewm(alpha=1/14, min_periods=14, adjust=False).mean()
            avg_loss = loss.ewm(alpha=1/14, min_periods=14, adjust=False).mean()
            rs = avg_gain / avg_loss
            rsi = 100 - (100 / (1 + rs))
            ema12 = close_s.ewm(span=12, adjust=False).mean()
            ema26 = close_s.ewm(span=26, adjust=False).mean()
            macd = ema12 - ema26
            signal = macd.ewm(span=9, adjust=False).mean()
            macd = macd - 0*signal  # 保留变量名一致；输出只需要macd
            mid = close_s.rolling(window=20, min_periods=20).mean()
            std = close_s.rolling(window=20, min_periods=20).std()
            boll_upper = mid + 2 * std
            boll_lower = mid - 2 * std
        
        # 组装数据
        df = stock_daily[["ts_code", "trade_date"]].copy()
        df["ma5"] = ma5.round(2)
        df["ma20"] = ma20.round(2)
        df["rsi"] = rsi.round(2)
        df["macd"] = macd.round(4)
        df["boll_upper"] = boll_upper.round(2)
        df["boll_lower"] = boll_lower.round(2)
        tech_data.append(df)
    
    # 合并数据并添加自增主键
    if tech_data:
        df_all = pd.concat(tech_data, ignore_index=True)
    else:
        df_all = pd.DataFrame(columns=["id","ts_code","trade_date","ma5","ma20","rsi","macd","boll_upper","boll_lower"])
        return df_all
    df_all["id"] = range(1, len(df_all)+1)
    # 调整字段顺序
    df_all = df_all[["id", "ts_code", "trade_date", "ma5", "ma20", "rsi", "macd", "boll_upper", "boll_lower"]]
    return df_all


# 6. 模拟AI预测结果（对应ai_predictions表，实际需接入模型）
def simulate_ai_predictions():
    predict_data = []
    predict_date = datetime.now().strftime("%Y%m%d")  # 预测执行日期（今天）
    for_date = (datetime.now() + timedelta(days=1)).strftime("%Y%m%d")  # 预测明天
    
    for i, ts_code in enumerate(stock_codes):
        # 模拟上涨概率（0-1之间，随机生成）
        prediction_score = round(0.5 + 0.3 * (i % 2), 2)  # 简单区分两支股票
        predict_data.append({
            "id": i + 1,
            "ts_code": ts_code,
            "predict_date": predict_date,
            "for_date": for_date,
            "prediction_score": prediction_score
        })
    
    return pd.DataFrame(predict_data)


# 7. 主函数：获取数据并保存为CSV
if __name__ == "__main__":
    # 获取各表数据
    stocks_df = get_stock_basic()
    daily_df = get_stock_daily()
    tech_df = get_technical_indicators(daily_df)
    predict_df = simulate_ai_predictions()
    
    # 保存为CSV（对应4张表）
    stocks_df.to_csv("stocks.csv", index=False)
    daily_df.to_csv("stock_daily_data.csv", index=False)
    tech_df.to_csv("stock_technical_indicators.csv", index=False)
    # 另存一份“清洗版”技术指标：去除含空值的行，便于直接做训练/验收
    tech_df_clean = tech_df.dropna(subset=["ma5", "ma20", "rsi", "macd", "boll_upper", "boll_lower"], how="any")
    tech_df_clean.to_csv("stock_technical_indicators_clean.csv", index=False)
    predict_df.to_csv("ai_predictions.csv", index=False)
    
    print("数据获取完成，已保存为CSV文件：")
    print("1. 股票基本信息表：stocks.csv")
    print("2. 日线行情表：stock_daily_data.csv")
    print("3. 技术指标表：stock_technical_indicators.csv")
    print("   技术指标清洗版：stock_technical_indicators_clean.csv（已去除空值行）")
    print("4. 预测结果表：ai_predictions.csv")